'''
Created on Jun 10, 2012
@author: oabalbin
Created on Jun 25, 2012
@author: alebalbin
'''

import logging
import argparse
import sys
import numpy as np
from collections import defaultdict
from lib.parse_gtf2 import read_gtf
from lib.samples_annotation import annotate_samples

JOB_SUCCESS=0
JOB_ERROR=1

class expMatrix:
    def __init__(self,n,m,
                 features_index,features_values,
                 features_types,
                 samples_dict):
        self.mat = np.zeros((n,m),dtype=np.float)
        self.mat_sense = np.zeros((n,m),dtype=np.float)
        self.mat_nat = np.zeros((n,m),dtype=np.float)
        self.samples=defaultdict()
        self.features_index = features_index
        self.features_values = features_values
        self.features_type=features_types
        self.features_exp_type=defaultdict(set)
        # Fill the matrix
        self.fill_expr_mat(samples_dict)#features_index,features_types, 
        
    def fill_expr_mat(self,samples_dict):#features_index,features_type,
        j=0
        for sp, fpath in samples_dict.iteritems():
            self.samples[j]=sp
            #self.read_cufflinks_fpkmtraking(j,features_index,
            #                       fpath)
            self.read_cufflinks_all_fpkmtraking(j,fpath) #features_index,features_type,
            
            j+=1

    def read_cufflinks_all_fpkmtraking(self,j,path_to_file):#features_index,features_type,
        ifile = open(path_to_file)
        hd = ifile.next()
        print path_to_file
        for l in ifile:
            f = l.strip('\n').split('\t')

            #feature=",".join(f[0:6]+f[7:13]).replace(',','_')
            #feature=",".join([f[0]]+[f[0]]+f[2:6]+[f[7]]+[f[7]]+f[9:13]).replace(',','_')
                
            feature=",".join([f[0]]+[f[0]]+[f[2].split('_')[0]]+f[3:6]+[f[7]]+[f[7]]+[f[9].split('_')[0]]+f[10:13]).replace(',','_')
            if len(f[2].split('_')) > 1:
                continue
    
            fpkm_snat_ratio=f[-1]
            fpkm_sense=f[6]
            fpkm_nat=f[13]
            #print feature,fpkm_snat_ratio
            i=self.features_index[feature]
            sense_type = self.features_type[feature]# features_type[feature].split('&')
            gene_biotype0,gene_biotype1=sense_type[3].split("gf.")[1],sense_type[4].split("gf.")[1]
            
            #print sense_type
            self.mat[i,j]=float(fpkm_snat_ratio)
            # One feature is a coding gene the other is not
            if sense_type[0]=="sense" and sense_type[1]=="antisense":
                #print feature,sense_type,fpkm_sense,fpkm_nat
                self.mat_sense[i,j]=float(fpkm_sense)
                self.mat_nat[i,j]=float(fpkm_nat)
                #map(self.features_type[feature].append,["sense","antisense"])
                self.features_exp_type[feature].add(("sense","antisense"))
            # One feature is a coding gene the other is not
            elif sense_type[0]=="antisense" and sense_type[1]=="sense":
                self.mat_sense[i,j]=float(fpkm_nat)
                self.mat_nat[i,j]=float(fpkm_sense)
                #map(self.features_type[feature].append,["antisense","sense"])
                self.features_exp_type[feature].add(("antisense","sense"))
            elif gene_biotype1=="antisense":
                self.mat_sense[i,j]=float(fpkm_sense)
                self.mat_nat[i,j]=float(fpkm_nat)
                #map(self.features_type[feature].append,["sense","antisense"])
                self.features_exp_type[feature].add(("sense","antisense"))
            elif gene_biotype0=="antisense":
                self.mat_sense[i,j]=float(fpkm_nat)
                self.mat_nat[i,j]=float(fpkm_sense)
                #map(self.features_type[feature].append,["antisense","sense"]) 
                self.features_exp_type[feature].add(("antisense","sense"))
            else:
                self.mat_sense[i,j]=float(fpkm_sense)
                self.mat_nat[i,j]=float(fpkm_nat)
                #map(self.features_type[feature].append,["posstrand","negstrand"])
                self.features_exp_type[feature].add(("posstrand","negstrand"))

            '''
            # Deprecated on 08-16-12: 
            # fpkm_sense > fpkm_nat could be positive and negative for the same pair in two different samples 
            # Both features are not protein coding genes
            # choose the one with bigger fpkm as the sense
            elif fpkm_sense > fpkm_nat:
                #print feature,sense_type,fpkm_sense,fpkm_nat
                self.mat_sense[i,j]=float(fpkm_sense)
                self.mat_nat[i,j]=float(fpkm_nat)
            '''

        ifile.close()

            
    def read_cufflinks_fpkmtraking(self,j,features_index,
                                   path_to_file):
        ifile = open(path_to_file)
        hd = ifile.next()
        print path_to_file
        for l in ifile:
            f = l.strip('\n').split('\t')
            feature=",".join(f[0:6]+f[7:13]).replace(',','_')
            fpkm_snat_ratio=f[-1]
            #print feature,fpkm_snat_ratio
            i=features_index[feature]
            self.mat[i,j]=float(fpkm_snat_ratio)        
        ifile.close()

    def read_cufflinks_sense_fpkmtraking(self,j,features_index,
                                   path_to_file):
        ifile = open(path_to_file)
        hd = ifile.next()
        print path_to_file
        for l in ifile:
            f = l.strip('\n').split('\t')
            feature=",".join(f[0:6]+f[7:13]).replace(',','_')
            fpkm_sense=f[6]
            #print feature,fpkm_snat_ratio
            i=features_index[feature]
            self.mat_sense[i,j]=float(fpkm_sense)        
        ifile.close()

    def read_cufflinks_nat_fpkmtraking(self,j,features_index,
                                   path_to_file):
        ifile = open(path_to_file)
        hd = ifile.next()
        print path_to_file
        for l in ifile:
            f = l.strip('\n').split('\t')
            feature=",".join(f[0:6]+f[7:13]).replace(',','_')
            fpkm_nat=f[13]
            #print feature,fpkm_snat_ratio
            i=features_index[feature]
            self.mat_nat[i,j]=float(fpkm_nat)        
        ifile.close()

   
    def read_nsat_fpkmtraking(self,j,features_index,
                                   path_to_file):
        '''
        what is the feature name
        ENST00000456328 ENSG00000223972 1       11869   14409   +       0.0103819338    ENST00000438504 ENSG00000227232 4       123966  125423  -       0.0428275534
        '''
        ifile = open(path_to_file)
        hd = ifile.next()
        for l in ifile:
            f = l.strip('\n').split('\t')
            feature,fpkm=f[0],f[9]
            
            print path_to_file
            #print feature
            i=features_index[feature]
            self.mat[i,j]=float(fpkm)        
        ifile.close()
        
                              
    def write_all_matrices(self,output_file,
                           annotation_file=None,mutations_file=None):
        output_file_ratio=output_file+"_ratio.txt"
        output_file_sense=output_file+"_sense.txt"
        output_file_nat=output_file+"_nat.txt"
        
        self.write_expr_matrix(self.mat,output_file_ratio,annotation_file,mutations_file)
        self.write_expr_matrix(self.mat_sense,output_file_sense,annotation_file,mutations_file)
        self.write_expr_matrix(self.mat_nat,output_file_nat,annotation_file,mutations_file)

        
    def write_expr_matrix(self,mat,output_file,
                          annotation_file=None,mutations_file=None):
        '''
        '''
        ofile=open(output_file,'w')
        hd=[]
        for j in range(len(self.samples)):
            hd.append(self.samples[j]) 
        #hd=['sgene_id','stranscript_id','schrom', 'sstart', 'send','rgene_id','rtranscript_id','rchrom', 'rstart', 'rend']
        if annotation_file is not None and mutations_file is not None:
            header=annotate_samples(annotation_file,mutations_file, self.samples)
            header.append(['nsat_pair']+hd)
        else:
            header=['nsat_pair']+hd
        for line in header:
            ofile.write(",".join(line).replace(',', '\t')+'\n')
        
        for i in range(mat.shape[0]):
            names=self.features_values[i]
            name_types=self.features_type[names]
            #print names,name_types,self.features_exp_type[names]
            # To handle pairs that were not seen in the transcriptome data
            # TODO: find a better way to do this.
            tmp=list(self.features_exp_type[names])
            if tmp:
                name_exp_types=list(tmp[0]) #list(list(a[1])[0])
            else:
                name_exp_types=["notseen","notseen"]
                print names
            
            names=names+"_"+",".join(name_types+name_exp_types).replace(',','_')
            ol = [names] + list(map(str,mat[i,:]))
            #print names
            ofile.write(",".join(ol).replace(',', '\t')+'\n')
        ofile.close()
        
    def write_expr_matrix_bytype(self,output_file,transcript_lookup,this_type):
        '''
        '''
        ofile=open(output_file,'w')
        hd=[]
        trackingid=0
        for j in range(len(self.samples)):
            hd.append(self.samples[j]) 
        hd = ['tracking_id','class_code','nearest_ref_id',
              'gene_id','gene_short_name','tss_id','locus','feature_type','strand']+\
              hd
        ofile.write(",".join(hd).replace(',', '\t')+'\n')
        
        for i in range(self.mat.shape[0]):
            names=self.features_values[i]
            names[trackingid]
            feature=transcript_lookup[names[trackingid]]
            names=[n.replace(',',';') for n in names]
            names.append(feature.feature_type)
            names.append(feature.strand)
            
            ol = names + list(map(str,self.mat[i,:]))
            ofile.write(",".join(ol).replace(',', '\t')+'\n')
        ofile.close()

class gene_lookup():
    def __init__(self,gtf_transcript):
        self.geneid=gtf_transcript.annotation["gene_id"]
        self.transcriptid=gtf_transcript.annotation["transcript_id"]
        self.feature_type=gtf_transcript.feature
        self.strand=gtf_transcript.strand

        
def count_lines(input_file):
    '''
    first line is header
    '''
    ifile = open(input_file)
    hd = ifile.next()
    features_index = defaultdict()
    features_values= defaultdict()
    for n,l in enumerate(ifile):
        f = l.strip('\n').split('\t')
        name=",".join(f[0:6]+f[7:13]).replace(',','_')
        if name == 'ENST00000499853_ENSG00000244932_3_129100141_129112992_+_ENSG00000251474_3_129101765_129115979_-':
            print name
            sys.exit(0)

        features_index[name]=n
        features_values[n]= name
    ifile.close()
    return n+1,features_index,features_values

def count_all_lines(input_files,features_sense,features_names):
    '''
    first line is header
    '''
    #hd = ifile.next()

    features_index = defaultdict()
    features_values= defaultdict()
    feature_types=defaultdict()
    unique_features=set()
    
    for s,p in input_files.iteritems():
        ifile = open(p)
        for n,l in enumerate(ifile):
            f = l.strip('\n').split('\t')
            name=",".join([f[0]]+[f[0]]+[f[2].split('_')[0]]+f[3:6]+[f[7]]+[f[7]]+[f[9].split('_')[0]]+f[10:13]).replace(',','_')# the f[2].split('_') is to correct for some errors with gene names such as ENST00000383621_HSCHR6_MHC_QBL_ that made me lost the chromosomal information
            if len(f[2].split('_')) > 1:
                continue

            f1,f2=f[0],f[7]
            f1_sense,f2_sense,pair_type,common_name1,common_name2=\
            features_sense[f1],features_sense[f2],f[14],features_names[f1],features_names[f2]
            #print name,n,f1_sense,f2_sense
            unique_features.add((name,0,f1_sense,f2_sense,pair_type,common_name1,common_name2))
        ifile.close()
    print "the number of features is", len(unique_features)
    #print unique_features
    for nind, feat in enumerate(unique_features):
        features_index[feat[0]]=nind#feat[1]
        features_values[nind]=feat[0]
        feature_types[feat[0]]=[feat[2],feat[3],feat[4],feat[5],feat[6]]
    print nind
    #print features_index
    return nind+1,features_index,features_values,feature_types

def read_gtf_ref(gtf_file):
    gtf_file =open(gtf_file)
    features_type=defaultdict()
    features_name=defaultdict()
    ft="transcript_id"
    ftn="transcript_name"
    ftb="gene_biotype"
    # It uses a gtf file template that should be the same
    # that was used for running cufflinks.
    
    # First step create transcripts joining exons that belong to the
    # same transcript.
    # It classifies the features in either sense or antisense
    # Being sense anything that is protein coding gene. 
    #gene_id "ENSG00000103152"; transcript_id "ENST00000436333"; exon_number "4"; gene_biotype "protein_coding"; gene_name "MPG"; p_id "P27734"; transcript_name "MPG-004"; tss_id "TSS29828";
    for transcript in read_gtf(gtf_file):
        transcript_annotation_name,transcript_feature, transcript_annotation=\
        transcript.annotation[ftn].replace('_','.'),\
        'gf.'+transcript.feature.replace('_','.'),'gb.'+transcript.annotation[ftb].replace('_','.')
        
        if transcript.feature =="protein_coding":
            gene_id=transcript.annotation[ft]
            features_type[gene_id]="sense"
            features_name[gene_id]=transcript_annotation_name+"_"+transcript_feature+"_"+transcript_annotation
        elif transcript.feature =="antisense":
            gene_id=transcript.annotation[ft]
            features_type[gene_id]="antisense"
            features_name[gene_id]=transcript_annotation_name+"_"+transcript_feature+"_"+transcript_annotation            
        else:
            gene_id=transcript.annotation[ft]
            features_type[gene_id]="antisense"
            features_name[gene_id]=transcript_annotation_name+"_"+transcript_feature+"_"+transcript_annotation
            
    gtf_file.close()
    return features_type,features_name
    
def read_samples_file(samples_file):
    '''
    '''
    myfiles=defaultdict()
    ifile=open(samples_file)
    for l in ifile:
        if l.startswith('#'):
            continue
        d=l.strip('\n').split('/')
        sample = d[7].split('.')[0] # in exds 6 this changes and it could be better to just fix the relation sample=file
        myfiles[sample]=l.strip('\n')
    ifile.close()
    return myfiles

    
def get_transcript_type(reference_gtf_file,this_type):
    '''
    this_type= either gene_id or transcript_id
    '''    
    features_lookup=defaultdict()
    
    gtf_file =open(reference_gtf_file)
    for transcript in read_gtf(gtf_file):
        if this_type=="gene_id":
            features_lookup[transcript.annotation["gene_id"]]=gene_lookup(transcript)
        elif this_type=="transcript_id":
            features_lookup[transcript.annotation["gene_id"]]=gene_lookup(transcript)
        else:
            continue
    
    gtf_file.close()
    return features_lookup


def main():
    '''
    This script uses the genes_tracking file to extract
    an expression matrix. 
    TODO: 08-16-2012: Needs to modify to include also transcripts predicted by cufflinks
    '''
    logging.basicConfig(level=logging.DEBUG,
                        format="%(asctime)s - %(name)s - %(levelname)s - %(message)s")
    logging.info("==========================")
    parser = argparse.ArgumentParser()
    parser.add_argument("--g", dest="gtf_file", default=None)
    parser.add_argument("--a", dest="annotation_file", default=None)
    parser.add_argument("--m", dest="mutation_file", default=None)
    parser.add_argument("samples_file")
    parser.add_argument("output_file")
    args = parser.parse_args()
    
    samples_dict = read_samples_file(args.samples_file)
    #print samples_dict.values()[0]
    #n,features_index,features_values = count_lines(samples_dict.values()[0])
    #print samples_dict
    genes_sense,genes_name = read_gtf_ref(args.gtf_file)
    print len(genes_sense)
    n,features_index,features_values,features_types = count_all_lines(samples_dict,genes_sense,genes_name)
    m=len(samples_dict)
    #print features_index
    print n,m
    cohort_matrix = expMatrix(n,m,
                              features_index,features_values,
                              features_types,
                              samples_dict)
    
    cohort_matrix.write_all_matrices(args.output_file,
                                     args.annotation_file,
                                     args.mutation_file)
    
    return JOB_SUCCESS

if __name__ == '__main__': 
    sys.exit(main())